"""
VLM文档处理器 - 使用Granite Docling视觉语言模型
"""

import os
import logging
import time
import psutil
import platform
from pathlib import Path
from typing import List, Dict, Any, Optional

from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import VlmPipelineOptions
from docling.datamodel.pipeline_options_vlm_model import InlineVlmOptions, InferenceFramework, ResponseFormat, TransformersModelType
from docling.pipeline.vlm_pipeline import VlmPipeline

# 尝试导入torch来检测MPS设备
try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

logger = logging.getLogger(__name__)


class VlmDoclingProcessor:
    """VLM文档处理器类 - 使用Granite Docling模型"""
    
    def __init__(self, model_manager=None):
        """初始化VLM处理器"""
        self.model_manager = model_manager
        self._converter = None
        self.max_file_size = 100 * 1024 * 1024  # 100MB

        # 支持的文件格式
        self.supported_formats = {
            'pdf', 'docx', 'pptx', 'xlsx', 'html', 'htm', 'md',
            'png', 'jpg', 'jpeg', 'tiff', 'bmp', 'gif', 'webp'
        }

        # 输出格式
        self.output_formats = ['markdown', 'json', 'html', 'text']

        # 记录VLM系统信息
        self._log_vlm_system_info()

        # 初始化转换器
        self._initialize_converter()

    def _log_vlm_system_info(self):
        """记录VLM系统信息和设备状态"""
        logger.info("=" * 60)
        logger.info("🧠 VLM处理器系统信息")
        logger.info("=" * 60)

        # 基本系统信息
        logger.info(f"操作系统: {platform.system()} {platform.release()}")
        logger.info(f"处理器: {platform.processor()}")

        # 内存信息
        memory = psutil.virtual_memory()
        logger.info(f"总内存: {memory.total / (1024**3):.1f} GB")
        logger.info(f"可用内存: {memory.available / (1024**3):.1f} GB")

        # VLM模型详细信息
        logger.info("🧠 VLM模型详细配置:")
        logger.info("   📚 核心模型: IBM Granite Docling 258M")
        logger.info("      - 模型类型: 视觉语言模型 (Vision-Language Model)")
        logger.info("      - 参数规模: 258M (2.58亿参数)")
        logger.info("      - 开发机构: IBM Research")
        logger.info("      - 模型架构: Transformer + Vision Encoder")
        logger.info("   🎯 模型能力:")
        logger.info("      🔍 视觉理解: 深度图像内容理解")
        logger.info("      📝 文档解析: 智能文档结构识别")
        logger.info("      📊 表格理解: 复杂表格结构解析")
        logger.info("      🧮 公式识别: 数学公式视觉识别")
        logger.info("      🔤 多语言OCR: 高精度文字识别")
        logger.info("      📋 布局分析: 智能页面布局理解")
        logger.info("   ⚙️  技术配置:")
        logger.info("      - 推理框架: Transformers (HuggingFace)")
        logger.info("      - 模型类型: AutoModel Vision2Seq")
        logger.info("      - 量化技术: 8-bit 量化加载")
        logger.info("      - 输出格式: Markdown结构化")
        logger.info("      - 最大令牌: 2048 tokens")
        logger.info("      - 温度参数: 0.0 (确定性输出)")

        # PyTorch和设备信息
        if TORCH_AVAILABLE:
            logger.info(f"PyTorch版本: {torch.__version__}")

            # 检查MPS设备
            if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🚀 MPS设备状态: 可用")
                logger.info("   - VLM模型将使用GPU加速")
                logger.info("   - Metal Performance Shaders优化")
                logger.info("   - 预期处理速度: 显著提升")
            else:
                logger.info("⚠️  MPS设备: 不可用，将使用CPU")
                logger.info("   - VLM模型处理速度较慢")
                logger.info("   - 建议使用较小的文档")
        else:
            logger.warning("PyTorch未安装，VLM功能可能受限")

        logger.info("=" * 60)
    
    def _initialize_converter(self):
        """初始化VLM转换器"""
        try:
            logger.info("🔧 开始初始化VLM转换器...")
            init_start = time.time()

            # 记录初始化前的内存状态
            memory_before = psutil.virtual_memory()
            logger.info(f"💾 初始化前内存使用: {memory_before.percent:.1f}% ({memory_before.used / (1024**3):.1f} GB)")

            logger.info("⚙️  配置IBM Granite VLM模型选项...")
            logger.info("   🏷️  模型标识:")
            logger.info("      - HuggingFace仓库: ibm-granite/granite-docling-258M")
            logger.info("      - 模型版本: 最新稳定版")
            logger.info("      - 许可证: Apache 2.0")
            logger.info("   🔧 技术参数:")
            logger.info("      - 推理框架: Transformers 4.x")
            logger.info("      - 模型类型: AutoModel Vision2Seq")
            logger.info("      - 输出格式: Markdown (结构化)")
            logger.info("      - 最大令牌数: 2048 tokens")
            logger.info("      - 温度参数: 0.0 (确定性输出)")
            logger.info("      - 8位量化: 启用 (内存优化)")
            logger.info("      - 信任远程代码: 启用")
            logger.info("   🎯 处理提示词:")
            logger.info("      - 任务: 文档页面转Markdown格式")
            logger.info("      - 要求: 保持所有文本、表格和结构")
            logger.info("      - 优化: 高质量结构化输出")

            # 配置VLM选项 - 使用Granite Docling模型
            vlm_options = InlineVlmOptions(
                repo_id="ibm-granite/granite-docling-258M",
                inference_framework=InferenceFramework.TRANSFORMERS,
                transformers_model_type=TransformersModelType.AUTOMODEL_VISION2SEQ,
                response_format=ResponseFormat.MARKDOWN,
                prompt="Convert this document page to markdown format, preserving all text, tables, and structure.",
                max_new_tokens=2048,
                temperature=0.0,
                load_in_8bit=True,
                trust_remote_code=True
            )

            logger.info("✅ VLM选项配置完成")
            
            logger.info("⚙️  配置VLM处理管道...")
            logger.info("   🖼️  图像预处理:")
            logger.info("      - 页面图片生成: 启用 (文档转图像)")
            logger.info("      - 图片缩放比例: 2.0x (高分辨率)")
            logger.info("      - 图像质量: 高清模式")
            logger.info("   🧠 VLM处理管道:")
            logger.info("      - 视觉编码器: 图像特征提取")
            logger.info("      - 语言模型: 文本生成和结构化")
            logger.info("      - 多模态融合: 视觉+语言理解")

            # 配置VLM管道选项
            pipeline_options = VlmPipelineOptions(
                # 启用页面图片生成（VLM需要）
                generate_page_images=True,
                # 图片缩放
                images_scale=2.0,
                # VLM选项
                vlm_options=vlm_options
            )

            logger.info("🔄 创建VLM文档转换器...")
            logger.info("   - 支持格式: PDF, DOCX, PPTX, XLSX, HTML, 图片, MD")
            logger.info("   - 使用VLM管道进行处理")

            # 创建文档转换器 - 使用VLM管道
            converter_start = time.time()
            self._converter = DocumentConverter(
                allowed_formats=[
                    InputFormat.PDF,
                    InputFormat.DOCX,
                    InputFormat.PPTX,
                    InputFormat.XLSX,
                    InputFormat.HTML,
                    InputFormat.IMAGE,
                    InputFormat.MD,
                ],
                format_options={
                    InputFormat.PDF: PdfFormatOption(
                        pipeline_cls=VlmPipeline,
                        pipeline_options=pipeline_options
                    ),
                    InputFormat.IMAGE: PdfFormatOption(
                        pipeline_cls=VlmPipeline,
                        pipeline_options=pipeline_options
                    )
                }
            )
            converter_time = time.time() - converter_start

            # 记录初始化后的内存状态
            memory_after = psutil.virtual_memory()
            memory_used = memory_after.used - memory_before.used
            init_time = time.time() - init_start

            logger.info("✅ VLM转换器初始化成功!")
            logger.info("=" * 50)
            logger.info(f"⏱️  总初始化时间: {init_time:.2f} 秒")
            logger.info(f"🔧 转换器创建时间: {converter_time:.2f} 秒")
            logger.info(f"💾 内存增加: {memory_used / (1024**2):.1f} MB")
            logger.info(f"💾 当前内存使用: {memory_after.percent:.1f}%")

            if TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🚀 MPS设备已就绪，VLM模型将使用GPU加速")
            else:
                logger.info("🖥️  将使用CPU进行VLM处理")

            logger.info("=" * 50)
            
        except Exception as e:
            logger.error(f"初始化VLM转换器失败: {e}")
            raise
    
    def is_supported_format(self, filename: str) -> bool:
        """检查文件格式是否支持"""
        if not filename:
            return False
        
        extension = filename.lower().split('.')[-1]
        return extension in self.supported_formats
    
    def process_document(self, file_path: Path, output_formats: List[str] = None, max_pages: int = 50) -> Dict[str, Any]:
        """
        处理文档
        
        Args:
            file_path: 文档文件路径
            output_formats: 输出格式列表
            max_pages: 最大处理页数
            
        Returns:
            处理结果字典
        """
        if output_formats is None:
            output_formats = ['markdown']
        
        try:
            # 开始处理日志
            start_time = time.time()
            file_size = file_path.stat().st_size

            logger.info("🧠 开始VLM文档处理")
            logger.info("=" * 60)
            logger.info(f"📄 文件名: {file_path.name}")
            logger.info(f"📁 文件路径: {file_path}")
            logger.info(f"📊 文件大小: {file_size / (1024*1024):.2f} MB")
            logger.info(f"📋 最大页数: {max_pages}")
            logger.info(f"🎯 输出格式: {', '.join(output_formats)}")
            logger.info(f"🧠 处理器类型: VLM (视觉语言模型)")

            # 记录处理前的系统状态
            memory_before = psutil.virtual_memory()
            logger.info(f"💾 处理前内存使用: {memory_before.percent:.1f}% ({memory_before.used / (1024**3):.1f} GB)")

            # 检查MPS设备状态
            if TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🚀 VLM模型使用MPS设备加速")
                logger.info("   - GPU加速的视觉语言模型推理")
                logger.info("   - 预期处理质量: 极高")
                logger.info("   - 预期处理速度: 中等 (质量优先)")
            else:
                logger.info("🖥️  VLM模型使用CPU处理")
                logger.info("   - 处理速度较慢，但质量保证")

            logger.info("⏳ 开始IBM Granite VLM处理...")
            logger.info("🔄 VLM处理阶段:")
            logger.info("   1️⃣  文档预处理: 转换为高分辨率图像")
            logger.info("   2️⃣  视觉编码: Granite视觉编码器分析图像")
            logger.info("   3️⃣  多模态理解: 视觉+语言模型融合")
            logger.info("   4️⃣  结构化生成: 智能Markdown输出")
            logger.info("   🧠 AI模型: IBM Granite Docling 258M")
            logger.info("   🎯 处理模式: 高精度视觉语言理解")

            conversion_start = time.time()

            # 转换文档
            logger.info("🚀 启动IBM Granite VLM推理...")
            result = self._converter.convert(file_path)

            conversion_time = time.time() - conversion_start
            logger.info(f"✅ IBM Granite VLM处理完成，耗时: {conversion_time:.2f} 秒")

            # 记录转换结果
            pages_count = len(result.document.pages) if hasattr(result.document, 'pages') else 0
            logger.info(f"📖 处理页数: {pages_count}")

            # 记录处理后的内存状态
            memory_after = psutil.virtual_memory()
            memory_used = memory_after.used - memory_before.used
            logger.info(f"💾 处理后内存使用: {memory_after.percent:.1f}% (增加: {memory_used / (1024**2):.1f} MB)")

            # IBM Granite VLM特有的处理统计
            if pages_count > 0:
                avg_time_per_page = conversion_time / pages_count
                logger.info(f"🧠 IBM Granite VLM处理分析:")
                logger.info(f"   📊 性能统计:")
                logger.info(f"      - 平均每页处理时间: {avg_time_per_page:.2f} 秒")
                logger.info(f"      - VLM推理速度: {1/avg_time_per_page:.2f} 页/分钟")
                logger.info(f"   🎯 处理质量:")
                logger.info(f"      - 视觉理解精度: 高精度 (258M参数)")
                logger.info(f"      - 文档结构保持: 优秀")
                logger.info(f"      - 多模态融合: 视觉+语言深度理解")
                logger.info(f"   🔍 模型能力:")
                logger.info(f"      - 复杂布局识别: ✅")
                logger.info(f"      - 表格结构理解: ✅")
                logger.info(f"      - 数学公式识别: ✅")
                logger.info(f"      - 多语言文字识别: ✅")

            # 准备输出结果
            logger.info("📝 开始生成输出格式...")
            outputs = {}

            # 生成不同格式的输出
            for format_name in output_formats:
                format_start = time.time()
                logger.info(f"   🔄 生成 {format_name.upper()} 格式...")

                if format_name == 'markdown':
                    content = result.document.export_to_markdown()
                    outputs['markdown'] = content
                    logger.info(f"   ✅ Markdown生成完成 ({len(content)} 字符)")
                elif format_name == 'json':
                    content = result.document.export_to_dict()
                    outputs['json'] = content
                    logger.info(f"   ✅ JSON生成完成")
                elif format_name == 'html':
                    content = result.document.export_to_html()
                    outputs['html'] = content
                    logger.info(f"   ✅ HTML生成完成 ({len(content)} 字符)")
                elif format_name == 'text':
                    content = result.document.export_to_text()
                    outputs['text'] = content
                    logger.info(f"   ✅ Text生成完成 ({len(content)} 字符)")

                format_time = time.time() - format_start
                logger.info(f"   ⏱️  {format_name.upper()}格式生成耗时: {format_time:.2f} 秒")

            # 获取文档信息
            doc_info = {
                'title': getattr(result.document, 'title', '未知标题'),
                'page_count': pages_count,
                'format': str(result.input.format) if hasattr(result, 'input') else '未知',
                'file_size': file_size,
                'processing_time': time.time() - start_time,
                'conversion_time': conversion_time,
                'device_used': 'MPS (GPU)' if (TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else 'CPU',
                'processor_type': 'VLM (视觉语言模型)'
            }

            total_time = time.time() - start_time

            # VLM处理完成总结
            logger.info("🎉 VLM文档处理完成!")
            logger.info("=" * 60)
            logger.info(f"📄 文件: {file_path.name}")
            logger.info(f"📊 总耗时: {total_time:.2f} 秒")
            logger.info(f"📖 处理页数: {pages_count}")
            logger.info(f"🧠 处理器: VLM (视觉语言模型)")
            logger.info(f"🖥️  设备: {doc_info['device_used']}")
            logger.info(f"💾 内存使用: {memory_used / (1024**2):.1f} MB")
            logger.info(f"🎯 输出格式: {', '.join(outputs.keys())}")
            logger.info(f"⭐ 处理质量: 高精度视觉理解")
            logger.info("=" * 60)
            
            return {
                'success': True,
                'outputs': outputs,
                'document_info': doc_info,
                'message': f'成功处理文档，共 {doc_info["page_count"]} 页'
            }
            
        except Exception as e:
            logger.error(f"VLM处理文档时出错: {e}")
            return {
                'success': False,
                'error': str(e),
                'message': f'处理文档时出错: {e}'
            }
    
    @property
    def is_initialized(self) -> bool:
        """检查转换器是否已初始化"""
        return self._converter is not None
    
    def get_info(self) -> Dict[str, Any]:
        """获取处理器信息"""
        return {
            'processor_type': 'VLM (Granite Docling)',
            'model': 'ibm-granite/granite-docling-258M',
            'supported_formats': list(self.supported_formats),
            'output_formats': self.output_formats,
            'max_file_size': self.max_file_size,
            'is_initialized': self.is_initialized
        }
